// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
//
// SPDX-License-Identifier: Apache-2.0

#include <cstdint>

#define REDUCE_OP PoolType::SUM
#define REDUCE_DIM ReduceDim::REDUCE_ROW

#define BCAST_LLKOP EltwiseBinaryType::ELWMUL
#define BCAST_DIM BroadcastType::COL

#include "compute_kernel_api/bcast.h"
#include "compute_kernel_api/eltwise_binary.h"
#include "compute_kernel_api/layernorm.h"
#include "ttnn/operations/normalization/kernel_util/compute/numeric.h"

ALWI void ACQ() { acquire_dst(); }
ALWI void REL() { release_dst(); }

namespace NAMESPACE {
void MAIN {
    namespace kutil = norm::kernel_util;
    namespace numeric = kutil::compute::numeric;
    namespace policies = kutil::compute::policies;

    uint32_t NCHt = get_arg_val<uint32_t>(0);
    constexpr uint32_t Wt = get_compile_time_arg_val(0);
    constexpr uint32_t blk = get_compile_time_arg_val(1);
    constexpr uint32_t do_gamma = get_compile_time_arg_val(2);
    constexpr uint32_t do_beta = get_compile_time_arg_val(3);
    constexpr bool FLOAT32_DTYPE = get_compile_time_arg_val(4) == 1;
    constexpr bool FLOAT32_REDUCTION = get_compile_time_arg_val(5) == 1;
    constexpr bool LEGACY_RSQRT = get_compile_time_arg_val(6) == 1;
    constexpr uint32_t one_over_W = get_compile_time_arg_val(7);

    constexpr uint32_t onetile = 1;
    // reserve one tile for zeros on cb_in2
    // TODO(AP): check that if DST is indeed zeroed by release_dst (and initially), we can use it as zeroes

    // Note that the entire W dimension must fit in the intermed0 CB for this kernel to be correct
    constexpr auto cb_scaler = tt::CBIndex::c_2;  // single tile generated by the reader
    constexpr auto cb_eps = tt::CBIndex::c_3;     // single tile generated by the reader
    constexpr auto cb_in = tt::CBIndex::c_0;      // input x or a for fused pre-add (x=a+b)
    constexpr auto cb_inb = tt::CBIndex::c_1;     // input b for fused pre-add
    constexpr auto cb_out = tt::CBIndex::c_16;    // output
    constexpr auto cb_gamma = tt::CBIndex::c_5;
    constexpr auto cb_beta = tt::CBIndex::c_6;
#if defined RMSNORM and not defined FUSE_PRE_ADD
    constexpr uint32_t cb_xmm = cb_in;  // x minus mean
#else
    constexpr uint32_t cb_xmm = tt::CBIndex::c_24;  // x minus mean
#endif
    constexpr auto cb_ex = tt::CBIndex::c_18;      // E[x]
    constexpr auto cb_ex2 = tt::CBIndex::c_19;     // E[(x-E[x])^2]
    constexpr auto cb_xmm2 = tt::CBIndex::c_20;    // xmm^2
    constexpr auto cb_ex2pe = tt::CBIndex::c_21;   // E[(x-E[x])^2]+eps
    constexpr auto cb_fusion = tt::CBIndex::c_22;  // stream gamma/beta
    constexpr auto scaler0 = 0;
#ifdef FUSE_PRE_ADD
#ifdef RMSNORM
    constexpr uint32_t cb_x = cb_xmm;
#else
    constexpr uint32_t cb_x = tt::CBIndex::c_23;
#endif
#else
    constexpr uint32_t cb_x = cb_in;
#endif

#ifdef FUSE_PRE_ADD
    binary_op_init_common(cb_in, cb_inb, cb_x);
#else
    binary_op_init_common(cb_in, cb_in, cb_xmm2);
#endif

    cb_wait_front(cb_scaler, 1);  // comes from the reader
    cb_wait_front(cb_eps, 1);     // comes from the reader

    constexpr int cb_im_or_out = (do_gamma | do_beta) ? cb_fusion : cb_out;

    for (uint32_t ncht = 0; ncht < NCHt; ncht++) {
        constexpr int onetile = 1;
        constexpr int dst0 = 0;

/*
 * X + Y
 */
#ifdef FUSE_PRE_ADD
        reconfig_data_format(cb_in, cb_inb);
        pack_reconfig_data_format(cb_x);
        add_tiles_init(cb_in, cb_inb);
        for (uint32_t wt = 0; wt < Wt; wt += blk) {
            ACQ();
            cb_wait_front(cb_in, blk);
            cb_wait_front(cb_inb, blk);
            cb_reserve_back(cb_x, blk);
            for (uint32_t j = 0; j < blk; j++) {
                add_tiles(cb_in, cb_inb, j, j, j);
                pack_tile(j, cb_x);
            }
            REL();
            cb_push_back(cb_x, blk);  // push the sum into the same buffer
            cb_pop_front(cb_in, blk);
            cb_pop_front(cb_inb, blk);
        }
#ifndef RMSNORM
        reconfig_data_format(cb_in, cb_x, cb_inb, cb_scaler);
#else
        reconfig_data_format(cb_in, cb_x, cb_inb, cb_x);
#endif
        // by the end of this loop we should end up with Wt tiles in cb_x
#else
#ifdef RMSNORM
        reconfig_data_format(cb_in, cb_in);
        pack_reconfig_data_format(cb_xmm2);
#endif
#endif

#ifndef RMSNORM
        // E[x]
        numeric::row_wise_mean<FLOAT32_REDUCTION>(cb_x, cb_scaler, cb_ex, one_over_W, Wt, blk);

        // x - E[x]
        reconfig_data_format(cb_x, cb_ex);
        cb_reserve_back(cb_xmm, Wt);
        sub_bcast_cols_init_short(cb_x, cb_ex);
        for (uint32_t wt = 0; wt < Wt; wt += blk) {
            ACQ();
            for (uint32_t wtr = 0; wtr < blk; wtr++) {
                sub_tiles_bcast_cols(cb_x, cb_ex, wt + wtr, 0, wtr);
                pack_tile(wtr, cb_xmm);
            }
            cb_push_back(cb_xmm, blk);
            REL();
        }
        cb_pop_front(cb_ex, 1);
        cb_pop_front(cb_x, Wt);

#ifndef FUSE_PRE_ADD
        reconfig_data_format_srca(cb_x, cb_xmm);
#endif
#endif

        /* (x - E[x])^2
         * compute temp = xmm*xmm = (x-E[x])^2
         */
        mul_tiles_init(cb_xmm, cb_xmm);
        for (uint32_t wt = 0; wt < Wt; wt += blk) {
            cb_wait_front(cb_xmm, wt + blk);
            cb_reserve_back(cb_xmm2, blk);
            ACQ();
            for (uint32_t wtr = 0; wtr < blk; wtr++) {
                mul_tiles(cb_xmm, cb_xmm, wt + wtr, wt + wtr, wtr);
                pack_tile(wtr, cb_xmm2);
            }
            cb_push_back(cb_xmm2, blk);
            REL();
        }

#if defined RMSNORM and not defined FUSED_PRE_ADD
        reconfig_data_format(cb_xmm, cb_xmm2, cb_xmm, cb_scaler);
#endif

        // Var[x]
        numeric::row_wise_mean<FLOAT32_REDUCTION, policies::PopInputPolicy::POP>(
            cb_xmm2, cb_scaler, cb_ex2, one_over_W, Wt, blk);

        // Var[x] + eps
        reconfig_data_format(cb_ex2, cb_eps);
        ACQ();
        add_tiles_init(cb_ex2, cb_eps);
        add_tiles(cb_ex2, cb_eps, 0, 0, dst0);

        cb_reserve_back(cb_ex2pe, 1);  // 1
        rsqrt_tile_init<LEGACY_RSQRT>();
        rsqrt_tile<LEGACY_RSQRT>(dst0);
        pack_reconfig_data_format(cb_ex2pe);
        pack_tile(dst0, cb_ex2pe);
        cb_push_back(cb_ex2pe, 1);
        REL();
        cb_pop_front(cb_ex2, 1);

        // (x-E[x]) / sqrt(Var[x] + eps) * gamma + beta
        cb_wait_front(cb_ex2pe, 1);
        for (uint32_t wt = 0; wt < Wt; wt += blk) {
            reconfig_data_format(cb_xmm, cb_ex2pe);
            if constexpr (do_gamma == 0 && do_beta == 0) {
                pack_reconfig_data_format(cb_out);
            } else {
                pack_reconfig_data_format(cb_fusion);
            }
            cb_reserve_back(cb_im_or_out, blk);
#if defined RMSNORM and not defined FUSE_PRE_ADD
            reconfig_data_format_srca(cb_fusion, cb_xmm);
#endif
            ACQ();
            mul_bcast_cols_init_short(cb_xmm, cb_ex2pe);
            for (uint32_t wtr = 0; wtr < blk; wtr++) {
                // cb_xmm[wt+wtr] since we pop Wt from cb_xmm after the entire loop
                mul_tiles_bcast_cols(cb_xmm, cb_ex2pe, wt + wtr, 0, wtr);  // tile *= 1/(sum(exp(x)))
                pack_tile(wtr, cb_im_or_out);  // pack either to intermediate (cb_fusion or out0)
            }
            cb_push_back(cb_im_or_out, blk);  // if no gamma/beta are provided, this will be passed on to the writer
            REL();

            if constexpr (!(do_gamma == 0 && do_beta == 0)) {
#if defined RMSNORM and not defined FUSE_PRE_ADD
                reconfig_data_format_srca(cb_xmm, cb_fusion);
#endif
            }
            if constexpr (do_gamma) {
                if constexpr (do_beta == 0) {
                    pack_reconfig_data_format(cb_out);
                }
                reconfig_data_format_srcb(cb_ex2pe, cb_gamma);
                ACQ();
                uint32_t cb_outg = do_beta ? cb_fusion : cb_out;
                mul_bcast_rows_init_short(cb_fusion, cb_gamma);
                cb_reserve_back(cb_outg, blk);
                cb_wait_front(cb_gamma, wt + blk);  // we don't pop, TODO: only wait on first ht
                cb_wait_front(cb_fusion, blk);
                for (uint32_t wtr = 0; wtr < blk; wtr++) {
                    mul_tiles_bcast_rows(cb_fusion, cb_gamma, wtr, wt + wtr, wtr);  // tile *= 1/(sum(exp(x)))
                    pack_tile(wtr, cb_outg);  // pack either to intermediate (cb_fusion or out0)
                }
                cb_pop_front(cb_fusion, blk);
                // we don't pop gamma
                cb_push_back(cb_outg, blk);
                // We don't pop gamma since it's 1,1,1,Wt and we reuse it for all NCHt
                REL();
            }
            if constexpr (do_beta) {
                pack_reconfig_data_format(cb_out);
                if constexpr (do_gamma) {
                    reconfig_data_format_srcb(cb_gamma, cb_beta);
                } else {
                    reconfig_data_format_srcb(cb_ex2pe, cb_beta);
                }
                ACQ();
                add_bcast_rows_init_short(cb_fusion, cb_beta);
                cb_reserve_back(cb_out, blk);
                cb_wait_front(cb_beta, wt + blk);  // TODO: optimization - only wait on first ht
                cb_wait_front(cb_fusion, blk);
                for (uint32_t wtr = 0; wtr < blk; wtr++) {
                    add_tiles_bcast_rows(cb_fusion, cb_beta, wtr, wt + wtr, wtr);  // tile *= 1/(sum(exp(x)))
                    pack_tile(wtr, cb_out);  // pack either to intermediate (cb_fusion or out0)
                }
                cb_pop_front(cb_fusion, blk);
                // We don't pop beta since it's 1,1,1,Wt and we reuse it for all NCHt
                cb_push_back(cb_out, blk);
                REL();
            }
        }
        cb_pop_front(cb_ex2pe, 1);
        cb_pop_front(cb_xmm, Wt);

    }  // NCHt loop
}
}  // namespace NAMESPACE
